'''------------------------------------------------------------------------------------------------
 $ @Author: YSQ
 $ @Date: 2021-12-22 15:06:25
 $ @LastEditTime: 2021-12-22 23:49:23
 $ @LastEditors: Please set LastEditors
 $ @Description: tracker类的父类，主要用于存储Tracker的信息
 $ @FilePath: \YSQTrackerEvaluationToolKit_v1.0\EvalToolKit\tracker.py
--------------------------------------------------------------------------------------------------'''
import sys
sys.path.append('./')
import numpy as np
from Tools.GetIou import getIouArray
from Tools.bboxformat import xywh2xyxy
from Tools.MTree import MTree
from tqdm import tqdm

class Tracker():
    def __init__(self,name,dir,color,anno) -> None:
        self.name = name
        self.dir = dir
        self.color = color
        self.anno = anno

        self.ACC = {}
        self.SUCC = {}
        self.ERR = {}
        self.IOU = {}

        self.failureFrame = {}

    def getACC(self):
        Acc = np.zeros(50)
        vNum = 0
        for vName in tqdm(self.anno.keys(),desc=MTree(1,'Get '+self.name+'\'s ACC')):
            vNum += 1
            partAcc = np.zeros(50)
            for i in range(self.ERR[vName].shape[0]):
                partAcc[:int(self.ERR[vName][i])+1] += 1
            partAcc /= self.ERR[vName].shape[0]
            partAcc = 1 - partAcc
            self.ACC[vName] = partAcc
            Acc += partAcc
        self.ACC['all'] = Acc/vNum


    def getSUCC(self):
        succ = np.zeros(100)
        vNum = 0
        for vName in tqdm(self.anno.keys(),desc=MTree(1,'Get '+self.name+'\'s SUCC')):
            vNum += 1
            partsucc = np.zeros(100)
            for i in range(self.IOU[vName].shape[0]):
                partsucc[:int(self.IOU[vName][i]*100) + 1] += 1
            partsucc /= self.IOU[vName].shape[0]
            self.SUCC[vName] = partsucc
            succ += partsucc
        self.SUCC['all'] = succ/vNum


    def getERR(self,gt):
        for vName in tqdm(gt.keys(),desc=MTree(1,'Get '+self.name+'\'s ERR')):
            partGT = gt[vName]
            partBox = self.anno[vName]
            # 计算中心位置
            partGT = [partGT[:, 0] + partGT[:, 2] / 2, partGT[:, 1] + partGT[:, 3] / 2]
            partGT = np.array(partGT)
            # 计算中心位置
            partBox = [partBox[:, 0] + partBox[:, 2] / 2, partBox[:, 1] + partBox[:, 3] / 2]
            partBox = np.array(partBox)
            #计算误差
            err = ((partBox[0] - partGT[0]) ** 2 + (partBox[1] - partGT[1]) ** 2) ** 0.5
            # 保存
            self.ERR[vName] = err

    def getIOU(self,gt):
        for vName in tqdm(gt.keys(),desc=MTree(1,'Get '+self.name+'\'s IOU')):
            partGT = xywh2xyxy(gt[vName])
            partBox = xywh2xyxy(self.anno[vName])
            #计算误差
            iou = getIouArray(partGT,partBox)
            # 保存
            self.IOU[vName] = iou

    def getFailureVideos(self,failureIOUThreshold,failureFrameNumThreshold,starFrameAndEndFrame):
        for vName in tqdm(self.IOU.keys(),desc=MTree(1,'Get '+self.name+'\'s FailureFrame')):
            failureFN = 0       # 存储跟踪失败的帧数
            frameList= []       # 存储跟踪失败帧的idx
            for frame in range(self.IOU[vName].shape[0]):
                if self.IOU[vName][frame]<failureIOUThreshold:
                    failureFN += 1
                    if vName in starFrameAndEndFrame.keys():
                        frameList.append(frame+starFrameAndEndFrame[vName][0]+1)
                    else:
                        frameList.append(frame+1)
            if failureFN>failureFrameNumThreshold:
                self.failureFrame[vName] = frameList

